from PIL import Image
import numpy as np
from skimage import data,filters,feature
import matplotlib.pyplot as plt
import skimage.morphology as sm

def read_img(img_path):
    # 读取图片
    img = Image.open(img_path)
    original_img_array = np.array(img)
    gray_img_array = np.array(img.convert('L'))

    return original_img_array, gray_img_array


def edge_detection(gray_img_array, kind, ksize = 3, sigma=3):
    # 边缘检测
    if kind == 1:
        edges = filters.sobel(gray_img_array)
    elif kind == 2:
        edges = filters.scharr(gray_img_array)
    elif kind == 3:
        edges = filters.prewitt(gray_img_array)
    elif kind == 4:
        edges = filters.roberts(gray_img_array)
    elif kind == 5:
        edges = filters.laplace(gray_img_array, ksize = ksize)
    elif kind == 6:
        edges = feature.canny(gray_img_array, sigma = sigma)

    return edges

def draw_edges(original_img_array, edges_array, threshold, ):
    edges_array = 1 * (edges_array > threshold)  # 灰度点大于指定阈值的保留下来（即为明显边界）
    pos = edges_array.nonzero()
    pos_num = len(pos[0])

    # 对于留下的边界点保留单一绿线颜色
    for i in range(pos_num):
        y = pos[1][i]
        x = pos[0][i]
        original_img_array[x][y][0] = 0
        original_img_array[x][y][1] = 255
        original_img_array[x][y][2] = 0

    image = Image.fromarray(np.uint8(original_img_array))
    image.show()



if __name__ == '__main__':
    # img_path = r'H:\临时存放地\formal_data\HER2\first_generation\train\10X\g3\3+_4801_10000000.jpg'
    img_path = r'H:\临时存放地\formal_data\HER2\first_generation\train\10X\g0\0_4710_10000001.jpg'
    # img_path = r'H:\临时存放地\formal_data\HER2\first_generation\train\10X\g1\1+_4526_10000002.jpg'

    original_img_array, gray_img_array = read_img(img_path)
    edges = edge_detection(gray_img_array, 6)
    # edges = sm.dilation(edges, sm.disk(2))
    draw_edges(original_img_array, edges, 0.1)